# export_raw_onnx160.py
import torch
from ultralytics.nn.tasks import DetectionModel

# 加载模型
ckpt = torch.load("D:/CodeCNN/yolov8-study/runs/detect/train31/weights/best31n.pt", map_location="cpu")
model = DetectionModel(ckpt["model"].yaml)
model.load_state_dict(ckpt["model"].float().state_dict())
model.fuse()
model.eval().cpu()

# 导出原始输出 (1, num_outputs, 2100)
dummy = torch.randn(1, 3, 320, 320)
torch.onnx.export(
    model,
    dummy,
    "best31n_raw-1.onnx",
    input_names=["images"],
    output_names=["output"],
    opset_version=12,
    export_params=True,
    do_constant_folding=True
)

print("✅ Raw model exported (should be ~6 MB)")